# 3d_servo_simulator.py
from vpython import *
import time


class VirtualArm3D:
    def __init__(self):
        """初始化3D虚拟奥特曼手臂"""
        self.scene = canvas(title="🤖 奥特曼AI控制系统", width=1000, height=700, background=color.gray(0.1))
        self.scene.camera.pos = vector(0, 0, 8)

        # 创建底座和躯干
        self.torso = box(pos=vector(0, 0, 0), size=vector(3, 5, 2), color=vector(0.8, 0.2, 0.2))
        self.base = box(pos=vector(0, -3, 0), size=vector(5, 1, 4), color=color.blue)

        # 创建左臂
        self.left_shoulder = sphere(pos=vector(-1.8, 1.5, 0), radius=0.3, color=color.red)
        self.left_upper_arm = cylinder(pos=self.left_shoulder.pos, axis=vector(-1, 0, 0),
                                       radius=0.25, length=1.5, color=color.yellow)
        self.left_elbow = sphere(pos=self.left_shoulder.pos + vector(-1.5, 0, 0), radius=0.25, color=color.red)
        self.left_forearm = cylinder(pos=self.left_elbow.pos, axis=vector(-1, 0, 0),
                                     radius=0.2, length=1.2, color=color.orange)

        # 创建右臂（对称）
        self.right_shoulder = sphere(pos=vector(1.8, 1.5, 0), radius=0.3, color=color.red)
        self.right_upper_arm = cylinder(pos=self.right_shoulder.pos, axis=vector(1, 0, 0),
                                        radius=0.25, length=1.5, color=color.yellow)
        self.right_elbow = sphere(pos=self.right_shoulder.pos + vector(1.5, 0, 0), radius=0.25, color=color.red)
        self.right_forearm = cylinder(pos=self.right_elbow.pos, axis=vector(1, 0, 0),
                                      radius=0.2, length=1.2, color=color.orange)

        # 状态显示
        self.status_label = label(pos=vector(0, -4, 0), text="等待指令...", height=20, color=color.green)
        self.instruction_label = label(pos=vector(0, 4, 0), text="说'拥抱'或'挥手'", height=16, color=color.white)

        # 角度状态
        self.left_angle = 0
        self.right_angle = 0

    def update_arm(self, left_angle=None, right_angle=None):
        """更新手臂姿态"""
        if left_angle is not None:
            self.left_angle = left_angle
            left_rad = radians(left_angle)
            self.left_upper_arm.axis = vector(-cos(left_rad), sin(left_rad), 0) * 1.5
            self.left_elbow.pos = self.left_shoulder.pos + self.left_upper_arm.axis
            self.left_forearm.pos = self.left_elbow.pos
            self.left_forearm.axis = vector(-cos(left_rad), sin(left_rad), 0) * 1.2

        if right_angle is not None:
            self.right_angle = right_angle
            right_rad = radians(right_angle)
            self.right_upper_arm.axis = vector(cos(right_rad), sin(right_rad), 0) * 1.5
            self.right_elbow.pos = self.right_shoulder.pos + self.right_upper_arm.axis
            self.right_forearm.pos = self.right_elbow.pos
            self.right_forearm.axis = vector(cos(right_rad), sin(right_rad), 0) * 1.2

    def perform_action(self, action_name):
        """执行动作序列"""
        self.status_label.text = f"执行: {action_name}"

        if action_name == "hug":
            self._perform_hug()
        elif action_name == "wave":
            self._perform_wave()
        elif action_name == "reset":
            self._reset()

        self.status_label.text = "就绪"

    def _perform_hug(self):
        """拥抱动作"""
        print("🤗 开始拥抱动作")
        for angle in [0, 30, 60, 90, 120, 150, 180]:
            self.update_arm(left_angle=angle, right_angle=angle)
            time.sleep(0.3)
        time.sleep(1)
        self._reset()

    def _perform_wave(self):
        """挥手动作"""
        print("👋 开始挥手动作")
        for _ in range(3):
            for angle in [0, 45, 90, 45, 0]:
                self.update_arm(right_angle=angle)
                time.sleep(0.2)
        self._reset()

    def _reset(self):
        """复位"""
        self.update_arm(left_angle=0, right_angle=0)


def main():
    """主测试函数"""
    print("=== 3D虚拟舵机模拟器 ===")
    print("指令列表:")
    print("  hug   - 模拟拥抱")
    print("  wave  - 模拟挥手")
    print("  reset - 复位")
    print("  exit  - 退出\n")

    arm = VirtualArm3D()

    while True:
        cmd = input("请输入指令: ").strip().lower()

        if cmd == 'hug':
            arm.perform_hug()
        elif cmd == 'wave':
            arm.perform_wave()
        elif cmd == 'reset':
            arm.reset()
            print("🔄 已复位")
        elif cmd == 'exit':
            break
        else:
            print("⚠️ 未知指令！请输入 hug/wave/reset/exit")


if __name__ == "__main__":
    main()